# 文件名: opencv卡尺工具.py (V9.1 - 调用修正版)
import cv2
import numpy as np
import math
import random

# 高斯拟合函数，用于亚像素边缘定位
def gaussian_fit(y1, y2, y3):
    denominator = y1 - 2 * y2 + y3
    if abs(denominator) < 1e-6: return 0.0
    return 0.5 * (y1 - y3) / denominator

# 直线边缘检测函数
def line_edge_detection(img, pt1, pt2,
                       caliper_count=30, search_width=60, projection_length=20,
                       edge_polarity='positive', show_roi=True):
    # 将输入点转换为numpy数组，创建可视化图像副本
    pt1, pt2 = np.array(pt1, dtype=np.float32), np.array(pt2, dtype=np.float32)
    direction_vec = pt2 - pt1; vis_img = img.copy()
    # 计算直线方向角度
    angle_rad = math.atan2(direction_vec[1], direction_vec[0]); angle_deg = np.degrees(angle_rad)
    # 在直线上均匀分布卡尺测量点
    caliper_centers = [pt1 + (i / (caliper_count + 1)) * direction_vec for i in range(1, caliper_count + 1)]
    edge_points = []  # 存储检测到的边缘点
    
    # 遍历每个卡尺测量点
    for center_point in caliper_centers:
        xc, yc = center_point
        roi_w = projection_length; roi_h = search_width
        padded_w = int(roi_w * 1.5); padded_h = int(roi_h * 1.5)
        # 创建旋转矩形区域
        rot_rect = ((xc, yc), (padded_w, padded_h), angle_deg)
        box = cv2.boxPoints(rot_rect)
        x_min, y_min = np.min(box, axis=0); x_max, y_max = np.max(box, axis=0)
        # 确保ROI在图像范围内
        x_min, y_min = int(max(0, x_min)), int(max(0, y_min)); x_max, y_max = int(min(img.shape[1], x_max)), int(min(img.shape[0], y_max))
        if (x_max - x_min) < 3 or (y_max - y_min) < 3: continue
        sub_img = img[y_min:y_max, x_min:x_max]
        local_center = (xc - x_min, yc - y_min)
        # 旋转图像使测量方向垂直
        rot_mat = cv2.getRotationMatrix2D(local_center, angle_deg, 1.0)
        rotated_sub_img = cv2.warpAffine(sub_img, rot_mat, (x_max - x_min, y_max - y_min), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT)
        # 提取ROI区域
        roi_x_start = int(rotated_sub_img.shape[1]/2 - roi_w/2); roi_y_start = int(rotated_sub_img.shape[0]/2 - roi_h/2)
        roi = rotated_sub_img[roi_y_start : roi_y_start + roi_h, roi_x_start : roi_x_start + roi_w]
        if roi.shape[0] < 3 or roi.shape[1] < 3: continue
        
        # 可视化ROI区域
        if show_roi:
            roi_corners_local = np.array([[roi_x_start, roi_y_start], [roi_x_start + roi_w, roi_y_start], [roi_x_start + roi_w, roi_y_start + roi_h], [roi_x_start, roi_y_start + roi_h]], dtype=np.float32)
            inv_rot_mat = cv2.invertAffineTransform(rot_mat)
            original_corners = cv2.transform(np.array([roi_corners_local]), inv_rot_mat)[0]
            original_corners[:, 0] += x_min; original_corners[:, 1] += y_min
            cv2.polylines(vis_img, [np.int32(original_corners)], True, (0, 200, 200), thickness=1)
        
        # 边缘检测处理
        gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY) if roi.ndim == 3 else roi
        # 垂直方向投影（平均灰度值）
        projection = np.mean(gray, axis=1, dtype=np.float64)
        gradient = np.gradient(projection)  # 计算梯度
        
        # 根据边缘极性选择梯度方向
        if edge_polarity == 'positive': grad_values = gradient
        elif edge_polarity == 'negative': grad_values = -gradient
        else: grad_values = np.abs(gradient)
        
        if len(grad_values) < 3: continue
        max_idx = np.argmax(grad_values)  # 找到梯度最大值位置
        gradient_threshold = 1.0
        if grad_values[max_idx] < gradient_threshold: continue
        
        # 亚像素边缘定位
        if 1 <= max_idx < len(grad_values) - 1:
            delta = gaussian_fit(grad_values[max_idx-1], grad_values[max_idx], grad_values[max_idx+1])
            if abs(delta) > 1.5: continue
            subpixel_edge_y_in_roi = max_idx + delta
            # 将边缘点坐标转换回原图像坐标系
            edge_pt_in_rotated = (roi_x_start + roi_w / 2.0, roi_y_start + subpixel_edge_y_in_roi)
            pt_homogeneous = np.array([[edge_pt_in_rotated[0]], [edge_pt_in_rotated[1]], [1]])
            inv_rot_mat = cv2.invertAffineTransform(rot_mat)
            original_pt_in_sub_img = np.dot(inv_rot_mat, pt_homogeneous)
            final_pt = (original_pt_in_sub_img[0,0] + x_min, original_pt_in_sub_img[1,0] + y_min)
            edge_points.append(final_pt)
    
    # 使用RANSAC拟合直线
    detected_line = None
    if len(edge_points) >= 2:
        edge_points_np = np.array(edge_points, dtype=np.float32)
        line = cv2.fitLine(edge_points_np, cv2.DIST_HUBER, 0, 0.01, 0.01)
        vx, vy, x0, y0 = line.flatten()
        t = (edge_points_np[:, 0] - x0) * vx + (edge_points_np[:, 1] - y0) * vy
        t_min, t_max = np.min(t), np.max(t)
        pt_start = (int(x0 + t_min * vx), int(y0 + t_min * vy)); pt_end = (int(x0 + t_max * vx), int(y0 + t_max * vy))
        detected_line = (pt_start, pt_end)
    
    return detected_line, edge_points, vis_img

# 通过三点计算圆的中心和半径
def _calculate_circle_from_3_points(p1, p2, p3):
    p1, p2, p3 = np.array(p1), np.array(p2), np.array(p3)
    D = 2 * (p1[0] * (p2[1] - p3[1]) + p2[0] * (p3[1] - p1[1]) + p3[0] * (p1[1] - p2[1]))
    if abs(D) < 1e-6: return None, None
    Ux = ((p1[0]**2 + p1[1]**2) * (p2[1] - p3[1]) + (p2[0]**2 + p2[1]**2) * (p3[1] - p1[1]) + (p3[0]**2 + p3[1]**2) * (p1[1] - p2[1])) / D
    Uy = ((p1[0]**2 + p1[1]**2) * (p3[0] - p2[0]) + (p2[0]**2 + p2[1]**2) * (p1[0] - p3[0]) + (p3[0]**2 + p3[1]**2) * (p2[0] - p1[0])) / D
    center = np.array([Ux, Uy])
    radius = np.linalg.norm(p1 - center)
    return center, radius

# 使用RANSAC算法拟合圆
def fit_circle_ransac(points, threshold=2.0, iterations=100):
    if len(points) < 3: return (None, None), []
    best_inliers = []  # 最佳内点集合
    best_model = (None, None)  # 最佳模型参数
    
    # RANSAC迭代
    for _ in range(iterations):
        sample_points = random.sample(points, 3)  # 随机选择三个点
        center, radius = _calculate_circle_from_3_points(*sample_points)
        if center is None: continue
        points_np = np.array(points)
        # 计算所有点到圆的距离
        distances = np.abs(np.linalg.norm(points_np - center, axis=1) - radius)
        current_inliers_indices = np.where(distances < threshold)[0]
        current_inliers = [points[i] for i in current_inliers_indices]
        # 更新最佳模型
        if len(current_inliers) > len(best_inliers):
            best_inliers = current_inliers
            best_model = (center, radius)
    
    # 使用内点重新计算圆参数
    final_center, final_radius = None, None
    if len(best_inliers) >= 5:
        inliers_np = np.array(best_inliers, dtype=np.float32)
        try:
            (cx, cy), final_radius = cv2.minEnclosingCircle(inliers_np)
            final_center = (cx, cy)
        except cv2.error: pass
    
    return (final_center, final_radius), best_inliers

# 圆形边缘检测函数
def caliper_circle_detection(img, center, init_radius,
                            caliper_count=36, search_width=40, projection_length=15,
                            edge_polarity='positive', show_roi=True,
                            ransac_threshold=2.0, ransac_iterations=100):
    center = np.array(center, dtype=np.float32); vis_img = img.copy(); edge_points = []
    # 在圆周上均匀分布卡尺测量点
    angles_deg = np.linspace(0, 360, caliper_count, endpoint=False)
    
    # 遍历每个径向角度
    for radial_angle_deg in angles_deg:
        radial_angle_rad = np.deg2rad(radial_angle_deg)
        direction_vec = np.array([np.cos(radial_angle_rad), np.sin(radial_angle_rad)])
        center_point = center + direction_vec * init_radius  # 计算卡尺中心点
        xc, yc = center_point
        tangent_angle_deg = radial_angle_deg + 90  # 切线方向（垂直于径向）
        
        roi_w = projection_length; roi_h = search_width
        padded_w = int(roi_w * 1.5); padded_h = int(roi_h * 1.5)
        rot_rect = ((xc, yc), (padded_w, padded_h), tangent_angle_deg)
        box = cv2.boxPoints(rot_rect)
        x_min, y_min = np.min(box, axis=0); x_max, y_max = np.max(box, axis=0)
        x_min, y_min = int(max(0, x_min)), int(max(0, y_min)); x_max, y_max = int(min(img.shape[1], x_max)), int(min(img.shape[0], y_max))
        if (x_max - x_min) < 3 or (y_max - y_min) < 3: continue
        
        sub_img = img[y_min:y_max, x_min:x_max]
        local_center = (xc - x_min, yc - y_min)
        # 旋转图像使测量方向垂直
        rot_mat = cv2.getRotationMatrix2D(local_center, tangent_angle_deg, 1.0)
        rotated_sub_img = cv2.warpAffine(sub_img, rot_mat, (x_max - x_min, y_max - y_min), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT)
        roi_x_start = int(rotated_sub_img.shape[1]/2 - roi_w/2); roi_y_start = int(rotated_sub_img.shape[0]/2 - roi_h/2)
        roi = rotated_sub_img[roi_y_start : roi_y_start + roi_h, roi_x_start : roi_x_start + roi_w]
        if roi.shape[0] < 3 or roi.shape[1] < 3: continue
        
        # 可视化ROI区域
        if show_roi:
            roi_corners_local = np.array([[roi_x_start, roi_y_start], [roi_x_start + roi_w, roi_y_start], [roi_x_start + roi_w, roi_y_start + roi_h], [roi_x_start, roi_y_start + roi_h]], dtype=np.float32)
            inv_rot_mat = cv2.invertAffineTransform(rot_mat)
            original_corners = cv2.transform(np.array([roi_corners_local]), inv_rot_mat)[0]
            original_corners[:, 0] += x_min; original_corners[:, 1] += y_min
            cv2.polylines(vis_img, [np.int32(original_corners)], True, (0, 200, 200), thickness=1)
        
        # 边缘检测处理（与直线检测相同）
        gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY) if roi.ndim == 3 else roi
        projection = np.mean(gray, axis=1, dtype=np.float64)
        gradient = np.gradient(projection)
        if edge_polarity == 'positive': grad_values = gradient
        elif edge_polarity == 'negative': grad_values = -gradient
        else: grad_values = np.abs(gradient)
        
        if len(grad_values) < 3: continue
        max_idx = np.argmax(grad_values)
        gradient_threshold = 1.0
        if grad_values[max_idx] < gradient_threshold: continue
        
        if 1 <= max_idx < len(grad_values) - 1:
            delta = gaussian_fit(grad_values[max_idx-1], grad_values[max_idx], grad_values[max_idx+1])
            if abs(delta) > 1.5: continue
            subpixel_edge_y_in_roi = max_idx + delta
            edge_pt_in_rotated = (roi_x_start + roi_w / 2.0, roi_y_start + subpixel_edge_y_in_roi)
            pt_homogeneous = np.array([[edge_pt_in_rotated[0]], [edge_pt_in_rotated[1]], [1]])
            inv_rot_mat = cv2.invertAffineTransform(rot_mat)
            original_pt_in_sub_img = np.dot(inv_rot_mat, pt_homogeneous)
            final_pt = (original_pt_in_sub_img[0,0] + x_min, original_pt_in_sub_img[1,0] + y_min)
            edge_points.append(final_pt)
    
    # 使用RANSAC拟合圆
    (optimized_center, optimized_radius), inlier_points = fit_circle_ransac(
        edge_points, threshold=ransac_threshold, iterations=ransac_iterations)
    
    return (optimized_center, optimized_radius), edge_points, inlier_points, vis_img

# ==================== V9.1 核心修改部分：修正 __main__ ====================
if __name__ == "__main__":
    # 加载图像
    img = cv2.imread('work_sheet_04.png')
    if img is None:
        print("错误: 无法加载图像")
        exit()
    
    # vis_img 将作为最终显示的画布
    vis_img = img.copy()

    # ========== 直线检测部分 ==========
    pt1 = (569, 93); pt2 = (691, 456)
    detected_line, line_pts, line_vis = line_edge_detection(
        img.copy(), pt1, pt2,
        caliper_count=15, search_width=80, projection_length=30,
        edge_polarity='positive', show_roi=True
    )
    if detected_line is not None:
        # 将直线检测的可视化结果（ROI）合并到最终画布上
        vis_img = cv2.addWeighted(vis_img, 0.5, line_vis, 0.5, 0)
        cv2.line(vis_img, detected_line[0], detected_line[1], (0, 0, 255), 2)  # 绘制拟合直线
        for pt in line_pts:
            cv2.drawMarker(vis_img, tuple(map(int, pt)), (255, 0, 0), cv2.MARKER_CROSS, 10, 2)  # 标记边缘点

    # ========== 圆形检测部分 ==========
    init_center = (485, 163); init_radius = 40
    
    # 修正：使用4个变量来接收4个返回值
    (final_center, final_radius), all_circle_pts, inlier_pts, circle_vis = caliper_circle_detection(
        img.copy(), init_center, init_radius,
        caliper_count=32, search_width=40, projection_length=20,
        edge_polarity='negative', show_roi=True,
        ransac_threshold=2.0, ransac_iterations=200
    )
    
    if final_center is not None:
        # 将圆形检测的可视化结果（ROI）合并到最终画布上
        vis_img = cv2.addWeighted(vis_img, 0.5, circle_vis, 0.5, 0)
        
        # 重新绘制直线检测结果，避免被圆形ROI覆盖
        if detected_line is not None:
             cv2.line(vis_img, detected_line[0], detected_line[1], (0, 0, 255), 2)
             for pt in line_pts:
                cv2.drawMarker(vis_img, tuple(map(int, pt)), (255, 0, 0), cv2.MARKER_CROSS, 10, 2)

        # 绘制最终拟合的圆 (红色)
        cv2.circle(vis_img, tuple(map(int, final_center)), int(final_radius), (0, 0, 255), 2)
        
        # 可视化RANSAC结果：区分内点和外点
        all_pts_set = {tuple(map(int, pt)) for pt in all_circle_pts}
        inlier_pts_set = {tuple(map(int, pt)) for pt in inlier_pts}
        outlier_pts_set = all_pts_set - inlier_pts_set
        
        # 绘制内点 (黄色十字)
        for pt in inlier_pts_set:
            cv2.drawMarker(vis_img, pt, (0, 255, 255), cv2.MARKER_CROSS, 10, 2)
            
        # 绘制外点 (紫色叉)
        for pt in outlier_pts_set:
            cv2.drawMarker(vis_img, pt, (255, 0, 255), cv2.MARKER_TILTED_CROSS, 10, 2)

        # 输出检测结果信息
        print(f"找到 {len(all_circle_pts)} 个边缘点，其中 {len(inlier_pts)} 个是内点。")
        print(f"圆心偏移: {np.linalg.norm(np.array(final_center)-np.array(init_center)):.2f}px")
        print(f"半径误差: {abs(final_radius-init_radius):.2f}px")

    # 显示最终结果
    cv2.imshow("检测结果", vis_img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()
